package com.flink.hbase;

import org.apache.flink.connector.file.sink.FileSink;
import org.apache.flink.core.fs.Path;
import org.apache.flink.formats.parquet.ParquetWriterFactory;
import org.apache.flink.formats.parquet.avro.ParquetAvroWriters;
import org.apache.flink.streaming.api.functions.sink.filesystem.BucketAssigner;
import org.apache.flink.streaming.api.functions.sink.filesystem.OutputFileConfig;
import org.apache.flink.streaming.api.functions.sink.filesystem.bucketassigners.SimpleVersionedStringSerializer;
import org.apache.flink.streaming.api.functions.sink.filesystem.rollingpolicies.DefaultRollingPolicy;
import org.apache.flink.streaming.api.functions.sink.filesystem.bucketassigners.SimpleVersionedSerializer;
import org.apache.parquet.hadoop.metadata.CompressionCodecName;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.time.Duration;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;

/**
 * Parquet 文件 Sink 工厂类
 * 负责创建和配置 FileSink，支持分区、压缩和滚动策略
 */
public class ParquetFileSink {

    private static final Logger LOG = LoggerFactory.getLogger(ParquetFileSink.class);

    /**
     * 创建 Parquet FileSink
     */
    public static FileSink<User> createSink(String outputPath, ParquetSinkConfig config) {
        LOG.info("创建 Parquet FileSink: outputPath={}, config={}", outputPath, config);

        // 创建 Parquet Writer Factory
        ParquetWriterFactory<User> writerFactory = ParquetAvroWriters.forReflectRecord(User.class);
        
        // 配置压缩
        if (config.getCompressionCodec() != null) {
            LOG.info("设置压缩格式: {}", config.getCompressionCodec());
        }

        // 创建 FileSink Builder
        FileSink.DefaultRowFormatBuilder<User> builder = FileSink
                .forRowFormat(new Path(outputPath), writerFactory)
                .withBucketAssigner(new UserBucketAssigner())
                .withRollingPolicy(createRollingPolicy(config))
                .withOutputFileConfig(createOutputFileConfig(config));

        return builder.build();
    }

    /**
     * 创建滚动策略
     */
    private static DefaultRollingPolicy createRollingPolicy(ParquetSinkConfig config) {
        DefaultRollingPolicy.PolicyBuilder policyBuilder = DefaultRollingPolicy.builder();

        // 文件大小
        if (config.getMaxFileSizeMB() > 0) {
            policyBuilder.withMaxPartSize(config.getMaxFileSizeMB() * 1024L * 1024L);
        }

        // 滚动时间间隔
        if (config.getRolloverIntervalMs() > 0) {
            policyBuilder.withRolloverInterval(Duration.ofMillis(config.getRolloverIntervalMs()));
        }

        // 不活跃时间
        if (config.getInactivityIntervalMs() > 0) {
            policyBuilder.withInactivityInterval(Duration.ofMillis(config.getInactivityIntervalMs()));
        }

        DefaultRollingPolicy policy = policyBuilder.build();
        LOG.info("滚动策略配置: maxFileSize={}MB, rolloverInterval={}ms, inactivityInterval={}ms",
                config.getMaxFileSizeMB(), config.getRolloverIntervalMs(), config.getInactivityIntervalMs());

        return policy;
    }

    /**
     * 创建输出文件配置
     */
    private static OutputFileConfig createOutputFileConfig(ParquetSinkConfig config) {
        OutputFileConfig.OutputFileConfigBuilder configBuilder = OutputFileConfig.builder();

        if (config.getFilePrefix() != null && !config.getFilePrefix().trim().isEmpty()) {
            configBuilder.withPartPrefix(config.getFilePrefix());
        }

        if (config.getFileSuffix() != null && !config.getFileSuffix().trim().isEmpty()) {
            configBuilder.withPartSuffix(config.getFileSuffix());
        }

        OutputFileConfig fileConfig = configBuilder.build();
        LOG.info("输出文件配置: prefix={}, suffix={}", config.getFilePrefix(), config.getFileSuffix());

        return fileConfig;
    }

    /**
     * 用户分桶分配器
     * 根据分区日期 (dt) 字段进行分桶
     */
    public static class UserBucketAssigner implements BucketAssigner<User, String> {

        private static final long serialVersionUID = 1L;
        private static final Logger LOG = LoggerFactory.getLogger(UserBucketAssigner.class);

        @Override
        public String getBucketId(User user, Context context) {
            if (user == null || user.getDt() == null) {
                String defaultBucket = "dt=unknown";
                LOG.warn("用户数据或分区日期为空，使用默认分桶: {}", defaultBucket);
                return defaultBucket;
            }

            // 使用 dt 字段进行分区
            String bucketId = "dt=" + user.getDt();
            LOG.debug("分配分桶: userId={}, bucketId={}", user.getUserId(), bucketId);
            return bucketId;
        }

        @Override
        public SimpleVersionedSerializer<String> getSerializer() {
            return SimpleVersionedStringSerializer.INSTANCE;
        }
    }

    /**
     * Parquet Sink 配置类
     */
    public static class ParquetSinkConfig {
        private int maxFileSizeMB = 128;                    // 最大文件大小（MB）
        private long rolloverIntervalMs = 60000;           // 滚动间隔（毫秒）
        private long inactivityIntervalMs = 30000;         // 不活跃间隔（毫秒）
        private String filePrefix = "user-data";           // 文件前缀
        private String fileSuffix = ".parquet";            // 文件后缀
        private CompressionCodecName compressionCodec = CompressionCodecName.SNAPPY; // 压缩格式

        public ParquetSinkConfig() {
        }

        // 静态工厂方法
        public static ParquetSinkConfig defaultConfig() {
            return new ParquetSinkConfig();
        }

        public static ParquetSinkConfig create(int maxFileSizeMB, 
                                             long rolloverIntervalMs, 
                                             long inactivityIntervalMs) {
            ParquetSinkConfig config = new ParquetSinkConfig();
            config.setMaxFileSizeMB(maxFileSizeMB);
            config.setRolloverIntervalMs(rolloverIntervalMs);
            config.setInactivityIntervalMs(inactivityIntervalMs);
            return config;
        }

        // Getters and Setters
        public int getMaxFileSizeMB() {
            return maxFileSizeMB;
        }

        public ParquetSinkConfig setMaxFileSizeMB(int maxFileSizeMB) {
            this.maxFileSizeMB = maxFileSizeMB;
            return this;
        }

        public long getRolloverIntervalMs() {
            return rolloverIntervalMs;
        }

        public ParquetSinkConfig setRolloverIntervalMs(long rolloverIntervalMs) {
            this.rolloverIntervalMs = rolloverIntervalMs;
            return this;
        }

        public long getInactivityIntervalMs() {
            return inactivityIntervalMs;
        }

        public ParquetSinkConfig setInactivityIntervalMs(long inactivityIntervalMs) {
            this.inactivityIntervalMs = inactivityIntervalMs;
            return this;
        }

        public String getFilePrefix() {
            return filePrefix;
        }

        public ParquetSinkConfig setFilePrefix(String filePrefix) {
            this.filePrefix = filePrefix;
            return this;
        }

        public String getFileSuffix() {
            return fileSuffix;
        }

        public ParquetSinkConfig setFileSuffix(String fileSuffix) {
            this.fileSuffix = fileSuffix;
            return this;
        }

        public CompressionCodecName getCompressionCodec() {
            return compressionCodec;
        }

        public ParquetSinkConfig setCompressionCodec(CompressionCodecName compressionCodec) {
            this.compressionCodec = compressionCodec;
            return this;
        }

        @Override
        public String toString() {
            return "ParquetSinkConfig{" +
                   "maxFileSizeMB=" + maxFileSizeMB +
                   ", rolloverIntervalMs=" + rolloverIntervalMs +
                   ", inactivityIntervalMs=" + inactivityIntervalMs +
                   ", filePrefix='" + filePrefix + '\'' +
                   ", fileSuffix='" + fileSuffix + '\'' +
                   ", compressionCodec=" + compressionCodec +
                   '}';
        }
    }

    /**
     * Parquet Sink 工具方法
     */
    public static class ParquetSinkUtils {

        /**
         * 根据参数创建配置
         */
        public static ParquetSinkConfig createConfigFromParameters(
                String maxFileSizeMB,
                String rolloverIntervalMs,
                String inactivityIntervalMs,
                String filePrefix,
                String fileSuffix,
                String compressionCodec) {

            ParquetSinkConfig config = new ParquetSinkConfig();

            if (maxFileSizeMB != null && !maxFileSizeMB.trim().isEmpty()) {
                try {
                    config.setMaxFileSizeMB(Integer.parseInt(maxFileSizeMB.trim()));
                } catch (NumberFormatException e) {
                    LOG.warn("无效的文件大小参数: {}, 使用默认值", maxFileSizeMB);
                }
            }

            if (rolloverIntervalMs != null && !rolloverIntervalMs.trim().isEmpty()) {
                try {
                    config.setRolloverIntervalMs(Long.parseLong(rolloverIntervalMs.trim()));
                } catch (NumberFormatException e) {
                    LOG.warn("无效的滚动间隔参数: {}, 使用默认值", rolloverIntervalMs);
                }
            }

            if (inactivityIntervalMs != null && !inactivityIntervalMs.trim().isEmpty()) {
                try {
                    config.setInactivityIntervalMs(Long.parseLong(inactivityIntervalMs.trim()));
                } catch (NumberFormatException e) {
                    LOG.warn("无效的不活跃间隔参数: {}, 使用默认值", inactivityIntervalMs);
                }
            }

            if (filePrefix != null && !filePrefix.trim().isEmpty()) {
                config.setFilePrefix(filePrefix.trim());
            }

            if (fileSuffix != null && !fileSuffix.trim().isEmpty()) {
                config.setFileSuffix(fileSuffix.trim());
            }

            if (compressionCodec != null && !compressionCodec.trim().isEmpty()) {
                try {
                    CompressionCodecName codec = CompressionCodecName.valueOf(compressionCodec.trim().toUpperCase());
                    config.setCompressionCodec(codec);
                } catch (IllegalArgumentException e) {
                    LOG.warn("不支持的压缩格式: {}, 使用默认值 SNAPPY", compressionCodec);
                }
            }

            return config;
        }

        /**
         * 验证输出路径
         */
        public static String validateOutputPath(String outputPath) {
            if (outputPath == null || outputPath.trim().isEmpty()) {
                throw new IllegalArgumentException("输出路径不能为空");
            }

            String cleaned = outputPath.trim();
            
            // 确保路径以 / 结尾
            if (!cleaned.endsWith("/")) {
                cleaned += "/";
            }

            return cleaned;
        }
    }
} 